#include "kernelpolynomialmethod.h"

using arma::cx_mat;

#include <mkl.h>

// This operator (use by simply appending () to a KPM type 
// variable) calculates the desity of states, DoS. The operator 
// returns	the object it self for easy access to result.
KernelPolynomialMethod &KernelPolynomialMethod::calcDoS(int id)
{
	// First get the Hamilton matrix.
	sp_cx_mat H = getHamiltonMatrix(id);
	
	// A ratio is set according to how challenging the calc. are!
	double lanczos_to_kpm_ratio = 0.;

	// Now if the bounds need calculating?
	if(spectra == automatic)
	{
		// A ratio is set according to how challenging the calc. are!
		// Note Lanczos is a bit faster!
		lanczos_to_kpm_ratio = 0.4*(double)N_states/(double)(N_poly*N_rand/2.);
		
		// A modified updateProgress is made (as this method should only go to ratio)!
		auto updateLaczosProgress = [&] (const double &progress_step) {
			updateProgress(progress_step*lanczos_to_kpm_ratio);
		};

		// We first tridiagonalize with Lanczos.
		LanczosMethod lm;
		lm.N_states = N_states;
		if(updateProgress)
			lm.updateProgress = updateLaczosProgress;

		lm(H,randomPhaseState(H.n_cols));

		// Now find all eigenvalues!
		long long result = LAPACKE_dsterf( LAPACK_COL_MAJOR, lm.alpha.memptr(), lm.beta.memptr() );

		if(result != 0)
			return *this;
		
		// The specra bounds are then taken.
		E_min = arma::min(lm.alpha);
		E_max = arma::max(lm.alpha);

		// Compensate for errors using tol_bounds !
	}
	
	// constants
	static const double pi = arma::datum::pi;

	// Now rescale the hamilton!
	double a = (E_max - E_min) / (2 - tol_bounds);
	double b = (E_max + E_min) / 2;
	sp_cx_mat bI = b*arma::speye<sp_cx_mat>(H.n_rows,H.n_cols);
	sp_cx_mat H_tilde = (H-bI)/a;

	
	// The update step is calculated! Note that updateProgress(step)
	// must be called N times = (1 - ratio)/step. This will give a full 
	// update of 1!
	step = (1. - lanczos_to_kpm_ratio)/(double)(N_poly*N_rand/2.);

	// Now we find some moments (or expansion coefficients).
	cx_mat mu_all(N_poly,N_rand);
	_Cilk_for(long r = 0; r < N_rand; r++)
	{
		// A weighted polymonial coefficient (to get the avg as a sum).
		mu_all.col(r) = momentSet(H_tilde) / N_rand;
	}

	// And the sum is taken!
	cx_vec mu = arma::sum(mu_all,1); // along the rows!

	// The jackson kernel.
	const auto jacksonKernal = [&] (cx_vec mu) 
	{
		double factor3 = tan(pi/(N_poly+1));

		_Cilk_for(long e = 0.; e < mu.n_elem; e++)
		{
			double n = (double)e;
			double factor1 = (N_poly-n-1)*cos(pi*n/(N_poly+1));
			double factor2 = sin(pi*n/(N_poly+1));

			mu(e) *= (factor1 + factor2/factor3)/(N_poly+1);
		}
		return mu;
	};
	
	// Rescale energies!
	const vec ones = arma::ones<vec>(x.n_rows);
	vec x_tilde = (arma::real(x)-ones*b)/a;
	// Then we reconstruct the DoS.
	vec DoS_tilde = reconstructDoS(mu, jacksonKernal,x_tilde);
	// And rescale!
	DoS = DoS_tilde/a;

	try_updateProgress();

	return *this;
}

KernelPolynomialMethod::KernelPolynomialMethod() :
	// A pointer to the used system. The system contains variables 
	// needed to construct elements (coupling strengths etc.).
	system(&System::Std),
	// A double containing the current bias (for bi-layered 
	// systems).
	bias(0.),
	// A double vector containing the current k-point.
	k(arma::zeros<vec>(3)),
	// A  complex double vector containing the sample points used in
	// the calculations. The size of this has little influence on the
	// operation time.
	x(arma::zeros<cx_vec>(1)),
	// An enum type containing weither the spectra bounds of the Unit 
	// are userdefined or automatic. This can be set automatic or 
	// userdefined. automatic uses the Lanczos method to find bounds.
	spectra(automatic),
	// Two doubles containing the spectra bounds of the Unit. This must be
	// set or the result wil become terribly wrong.
	E_min(-10), E_max(10),
	// A double containing the tolerance of the bounds, as not all elements
	// are well bounded within E_min, E_max.
	tol_bounds(0.001),
	// An integer containing the numer of random-phase states to
	// use in the calculations.
	N_rand(1),
	// An integer containing the numer of polynomials to use in the
	// expansion in the calculations.
	N_poly(1000),
	// An integer containing the number of lanczos states to use in the
	// calculations of the bounds.
	N_states(1000)
{
}

// This function attempts to update the progress.
void KernelPolynomialMethod::try_updateProgress() const
{
	if(!updateProgress)
		return;

	updateProgress(step);
}

// This function returns a random phase state, i.e |r> = 1/R sum_1^R exp(i phi) |i>.
cx_vec KernelPolynomialMethod::randomPhaseState(size_t n)
{
	// constants
	static const complex i = complex(0.,1.);
	static const complex pi = arma::datum::pi;
	
	cx_vec phases = 2.*pi*i*arma::randu<vec>(n); // Random phases!
	return arma::exp(phases)/sqrt((double)n); // return normalized!
}
	
// This function returns the Hamilton matrix H.
sp_cx_mat KernelPolynomialMethod::getHamiltonMatrix(int id) const
{
	// Make a Hamiton object and set the options.
	Hamiltonian hamilton;
	hamilton.system = system;
	hamilton.bias = bias;
	hamilton.k = k;

	// Calculate and return the sparse matrix.
	return hamilton(id).H;
}

// This function returns the Chebyshev polynomial expansion coefficients mu_n (or moments)
// using a random phase as input.
cx_vec KernelPolynomialMethod::momentSet(const sp_cx_mat &H_tilde) const 
{ 
	// Simply run the function with a new random phase state.
	return momentSet(H_tilde, randomPhaseState(H_tilde.n_cols)); 
}

// This function returns the Chebyshev polynomial expansion coefficients mu_n (or moments)
// beginning from a specific state |r>
cx_vec KernelPolynomialMethod::momentSet(const sp_cx_mat &H_tilde, const cx_vec &r_state) const
{
	// The result is allocated to hold the moments.
	cx_vec mu(N_poly);

	// Note that the coefficients are found two at a time, due to simplifications.
	// Now we need three temporary states: the current, the previous, and the next.
	cx_vec r_p; // The previous
	cx_vec r_c = r_state; // The current
	cx_vec r_n = H_tilde*r_state; // The next
	
	// The initial coefficients are found:
	mu(0) = arma::cdot(r_c,r_c);
	mu(1) = arma::cdot(r_n,r_c);
	// Now the following N_Poly-2 coefficients are found!
	for(long n = 2; n < N_poly; n += 2)
	{
		// Now shift the temporary states.
		r_p = r_c; // prev = current!
		r_c = r_n; // current = next;

		// Get the next state from current and prev!
		r_n = 2*H_tilde*r_c - r_p; // Next = 2*H*current - prev!

		// Set the two next moments accordingly.
		mu(n)	= 2.*arma::cdot(r_c,r_c) - mu(0);
		mu(n+1) = 2.*arma::cdot(r_n,r_c) - mu(1);

		try_updateProgress();
	}

	// Return the coefficients.
	return mu;
}

// Reconstructs the element using a kernel.
vec KernelPolynomialMethod::reconstructDoS(cx_vec mu, const std::function< cx_vec (cx_vec)> &kernal, const vec &x_tilde) const
{
	// constants
	static const double pi = arma::datum::pi;

	// First we apply the kernel.
	mu = kernal(mu);

	vec DoS_tilde(x_tilde.n_rows);
	
	_Cilk_for(long nE = 0; nE < x_tilde.n_rows; nE++)
	{
		// Three temporary constants are created: previus, current, and next.
		double Tp;	// Prev.
		double Tc = x_tilde(nE); // current.
		double Tn = 1; // next.

		// Initial factor is added to the sum.
		DoS_tilde(nE) = mu(0).real();

		for(long n = 1; n < mu.n_rows; n++)
		{
			// The temps are shifted.
			Tp = Tc;
			Tc = Tn;

			// The next constant is found.
			Tn = 2*x_tilde(nE)*Tc - Tp;
			// the value is added to the sum.
			DoS_tilde(nE) += 2*mu(n).real()*Tn;
		}

		// The final value is scaled.
		DoS_tilde(nE) /= (pi*sqrt(1-pow(x_tilde(nE),2)));
	}

	return DoS_tilde;
}